{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Multi-Agent Orchestration with NeMo Agent Toolkit\n",
    "\n",
    "In this notebook, we showcase how the NVIDIA NeMo Agent Toolkit (NAT) can be used to use a mixture of inbuilt tools and agents, as well as custom tools and workflows. Multi-agent orchestration is an important concept in many agentic AI facets: notably runtime and token efficiency. Here we aim to show how simple multi-agent orchestration is able to be implemented when using NAT. We show how an orchestration agent can call tools and sub-agents to facilitate complex tasks.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "\n",
    "- [0.0) Setup](#setup)\n",
    "  - [0.1) Prerequisites](#prereqs)\n",
    "  - [0.2) API Keys](#api-keys)\n",
    "  - [0.3) Data Sources](#data-sources)\n",
    "  - [0.4) Installing NeMo Agent Toolkit](#installing-nat)\n",
    "- [1.0) Defining Tools for the Multi-Agent Workflow](#creating-workflow)\n",
    "  - [1.1) Total Product Sales Tool](#product-sales-tool)\n",
    "  - [1.2) Sales Per Day Tool](#sales-per-day-tool)\n",
    "  - [1.3) Detect Outliers Tool](#detect-outliers-tool)\n",
    "  - [1.4) Technical Specs Retrieval Tool](#technical-specs-tool)\n",
    "  - [1.5) Data Analysis/Plotting Tools](#plotting-tools)\n",
    "  - [1.6) Register The Tools](#register-tools)\n",
    "- [2.0) Adding an Agent Orchestrator](#adding-orchestrator)\n",
    "  - [2.1) Agent Orchestrator Workflow Configuration File](#orchestrator-config)\n",
    "  - [2.2) Running the Workflow](#running-orchestrator-workflow)\n",
    "- [3.0) Adding a Custom Agent](#adding-custom-agent)\n",
    "  - [3.1) Human-in-the-Loop (HITL) Approval Tool](#hitl-tool)\n",
    "  - [3.2) Graph Summarizer Tool](#graph-summarizer-tool)\n",
    "  - [3.3) Custom Data Visualization Agent With HITL Approval](#custom-viz-agent)\n",
    "  - [3.4) Custom Agent Workflow Configuration File](#custom-agent-config)\n",
    "  - [3.5) Running the Workflow](#running-custom-workflow)\n",
    "- [4.0) Next Steps](#next-steps)\n",
    "\n",
    "<span style=\"color:rgb(0, 31, 153); font-style: italic;\">Note: In Google Colab use the Table of Contents tab to navigate.</span>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By the conclusion of this example, we will create a simple mixture-of-agents that serves as an assistant in retail sales.\n",
    "\n",
    "> **Note:** _This is just an example agent system that uses dummy data. The intention is to demonstrate some of the capabilities of this toolkit and how a new user can get familiar with it._\n",
    "\n",
    "This agent system has:\n",
    "\n",
    "1. A **supervisor** agent that routes incoming requests to the downstream agent expert\n",
    "2. A **data insight agent** that is a tool-calling agent capable of answering questions about sales data\n",
    "3. A **RAG agent** that is capable of answering questions about products using context from a product catalog\n",
    "4. A **data visualization agent** that is capable of plotting graphs and trends\n",
    "\n",
    "We demonstrate the following capabilities:\n",
    "* RAG\n",
    "* Multi-framework support\n",
    "* Human-in-the-Loop\n",
    "* Multi-agent support\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"setup\"></a>\n",
    "# 0.0) Setup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"prereqs\"></a>\n",
    "## 0.1) Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before users run this notebook, it is encouraged that they run `3_adding_tools_to_agents.py`. As this notebook will assume that users are familiar with tool calling agents, and gloss over the details of the initial agent that has tool calling setup.\n",
    "\n",
    "- **Platform:** Linux, macOS, or Windows\n",
    "- **Python:** version 3.11, 3.12, or 3.13\n",
    "- **Python Packages:** `pip`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"api-keys\"></a>\n",
    "## 0.2) API Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this notebook, you will need the following API keys to run all examples end-to-end:\n",
    "\n",
    "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n",
    "\n",
    "Then you can run the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"data-sources\"></a>\n",
    "## 0.3) Data Sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Several data files are required for this example. To keep this as a stand-alone example, the files are included here as cells which can be run to create them.\n",
    "\n",
    "The following cell creates the `data` directory as well as a `rag` subdirectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data/rag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell writes the `data/retail_sales_data.csv` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile data/retail_sales_data.csv\n",
    "Date,StoreID,Product,UnitsSold,Revenue,Promotion\n",
    "2024-01-01,S001,Laptop,1,1000,No\n",
    "2024-01-01,S001,Phone,9,4500,No\n",
    "2024-01-01,S001,Tablet,2,600,No\n",
    "2024-01-01,S002,Laptop,9,9000,No\n",
    "2024-01-01,S002,Phone,10,5000,No\n",
    "2024-01-01,S002,Tablet,5,1500,No\n",
    "2024-01-02,S001,Laptop,4,4000,No\n",
    "2024-01-02,S001,Phone,11,5500,No\n",
    "2024-01-02,S001,Tablet,7,2100,No\n",
    "2024-01-02,S002,Laptop,7,7000,No\n",
    "2024-01-02,S002,Phone,6,3000,No\n",
    "2024-01-02,S002,Tablet,9,2700,No\n",
    "2024-01-03,S001,Laptop,6,6000,No\n",
    "2024-01-03,S001,Phone,7,3500,No\n",
    "2024-01-03,S001,Tablet,8,2400,No\n",
    "2024-01-03,S002,Laptop,3,3000,No\n",
    "2024-01-03,S002,Phone,16,8000,No\n",
    "2024-01-03,S002,Tablet,5,1500,No\n",
    "2024-01-04,S001,Laptop,5,5000,No\n",
    "2024-01-04,S001,Phone,11,5500,No\n",
    "2024-01-04,S001,Tablet,9,2700,No\n",
    "2024-01-04,S002,Laptop,2,2000,No\n",
    "2024-01-04,S002,Phone,12,6000,No\n",
    "2024-01-04,S002,Tablet,7,2100,No\n",
    "2024-01-05,S001,Laptop,8,8000,No\n",
    "2024-01-05,S001,Phone,18,9000,No\n",
    "2024-01-05,S001,Tablet,5,1500,No\n",
    "2024-01-05,S002,Laptop,7,7000,No\n",
    "2024-01-05,S002,Phone,10,5000,No\n",
    "2024-01-05,S002,Tablet,10,3000,No\n",
    "2024-01-06,S001,Laptop,9,9000,No\n",
    "2024-01-06,S001,Phone,11,5500,No\n",
    "2024-01-06,S001,Tablet,5,1500,No\n",
    "2024-01-06,S002,Laptop,5,5000,No\n",
    "2024-01-06,S002,Phone,14,7000,No\n",
    "2024-01-06,S002,Tablet,10,3000,No\n",
    "2024-01-07,S001,Laptop,2,2000,No\n",
    "2024-01-07,S001,Phone,15,7500,No\n",
    "2024-01-07,S001,Tablet,6,1800,No\n",
    "2024-01-07,S002,Laptop,0,0,No\n",
    "2024-01-07,S002,Phone,7,3500,No\n",
    "2024-01-07,S002,Tablet,12,3600,No\n",
    "2024-01-08,S001,Laptop,5,5000,No\n",
    "2024-01-08,S001,Phone,8,4000,No\n",
    "2024-01-08,S001,Tablet,5,1500,No\n",
    "2024-01-08,S002,Laptop,4,4000,No\n",
    "2024-01-08,S002,Phone,11,5500,No\n",
    "2024-01-08,S002,Tablet,9,2700,No\n",
    "2024-01-09,S001,Laptop,6,6000,No\n",
    "2024-01-09,S001,Phone,9,4500,No\n",
    "2024-01-09,S001,Tablet,8,2400,No\n",
    "2024-01-09,S002,Laptop,7,7000,No\n",
    "2024-01-09,S002,Phone,11,5500,No\n",
    "2024-01-09,S002,Tablet,8,2400,No\n",
    "2024-01-10,S001,Laptop,6,6000,No\n",
    "2024-01-10,S001,Phone,11,5500,No\n",
    "2024-01-10,S001,Tablet,5,1500,No\n",
    "2024-01-10,S002,Laptop,8,8000,No\n",
    "2024-01-10,S002,Phone,5,2500,No\n",
    "2024-01-10,S002,Tablet,6,1800,No\n",
    "2024-01-11,S001,Laptop,5,5000,No\n",
    "2024-01-11,S001,Phone,7,3500,No\n",
    "2024-01-11,S001,Tablet,5,1500,No\n",
    "2024-01-11,S002,Laptop,4,4000,No\n",
    "2024-01-11,S002,Phone,10,5000,No\n",
    "2024-01-11,S002,Tablet,4,1200,No\n",
    "2024-01-12,S001,Laptop,2,2000,No\n",
    "2024-01-12,S001,Phone,10,5000,No\n",
    "2024-01-12,S001,Tablet,9,2700,No\n",
    "2024-01-12,S002,Laptop,8,8000,No\n",
    "2024-01-12,S002,Phone,10,5000,No\n",
    "2024-01-12,S002,Tablet,14,4200,No\n",
    "2024-01-13,S001,Laptop,3,3000,No\n",
    "2024-01-13,S001,Phone,6,3000,No\n",
    "2024-01-13,S001,Tablet,9,2700,No\n",
    "2024-01-13,S002,Laptop,1,1000,No\n",
    "2024-01-13,S002,Phone,12,6000,No\n",
    "2024-01-13,S002,Tablet,7,2100,No\n",
    "2024-01-14,S001,Laptop,4,4000,Yes\n",
    "2024-01-14,S001,Phone,16,8000,Yes\n",
    "2024-01-14,S001,Tablet,4,1200,Yes\n",
    "2024-01-14,S002,Laptop,5,5000,Yes\n",
    "2024-01-14,S002,Phone,14,7000,Yes\n",
    "2024-01-14,S002,Tablet,6,1800,Yes\n",
    "2024-01-15,S001,Laptop,9,9000,No\n",
    "2024-01-15,S001,Phone,6,3000,No\n",
    "2024-01-15,S001,Tablet,11,3300,No\n",
    "2024-01-15,S002,Laptop,5,5000,No\n",
    "2024-01-15,S002,Phone,10,5000,No\n",
    "2024-01-15,S002,Tablet,4,1200,No\n",
    "2024-01-16,S001,Laptop,6,6000,No\n",
    "2024-01-16,S001,Phone,11,5500,No\n",
    "2024-01-16,S001,Tablet,5,1500,No\n",
    "2024-01-16,S002,Laptop,4,4000,No\n",
    "2024-01-16,S002,Phone,7,3500,No\n",
    "2024-01-16,S002,Tablet,4,1200,No\n",
    "2024-01-17,S001,Laptop,6,6000,No\n",
    "2024-01-17,S001,Phone,14,7000,No\n",
    "2024-01-17,S001,Tablet,7,2100,No\n",
    "2024-01-17,S002,Laptop,3,3000,No\n",
    "2024-01-17,S002,Phone,7,3500,No\n",
    "2024-01-17,S002,Tablet,6,1800,No\n",
    "2024-01-18,S001,Laptop,7,7000,Yes\n",
    "2024-01-18,S001,Phone,10,5000,Yes\n",
    "2024-01-18,S001,Tablet,6,1800,Yes\n",
    "2024-01-18,S002,Laptop,5,5000,Yes\n",
    "2024-01-18,S002,Phone,16,8000,Yes\n",
    "2024-01-18,S002,Tablet,8,2400,Yes\n",
    "2024-01-19,S001,Laptop,4,4000,No\n",
    "2024-01-19,S001,Phone,12,6000,No\n",
    "2024-01-19,S001,Tablet,7,2100,No\n",
    "2024-01-19,S002,Laptop,3,3000,No\n",
    "2024-01-19,S002,Phone,12,6000,No\n",
    "2024-01-19,S002,Tablet,8,2400,No\n",
    "2024-01-20,S001,Laptop,6,6000,No\n",
    "2024-01-20,S001,Phone,8,4000,No\n",
    "2024-01-20,S001,Tablet,6,1800,No\n",
    "2024-01-20,S002,Laptop,8,8000,No\n",
    "2024-01-20,S002,Phone,9,4500,No\n",
    "2024-01-20,S002,Tablet,8,2400,No\n",
    "2024-01-21,S001,Laptop,3,3000,No\n",
    "2024-01-21,S001,Phone,9,4500,No\n",
    "2024-01-21,S001,Tablet,5,1500,No\n",
    "2024-01-21,S002,Laptop,8,8000,No\n",
    "2024-01-21,S002,Phone,15,7500,No\n",
    "2024-01-21,S002,Tablet,7,2100,No\n",
    "2024-01-22,S001,Laptop,1,1000,No\n",
    "2024-01-22,S001,Phone,15,7500,No\n",
    "2024-01-22,S001,Tablet,5,1500,No\n",
    "2024-01-22,S002,Laptop,11,11000,No\n",
    "2024-01-22,S002,Phone,4,2000,No\n",
    "2024-01-22,S002,Tablet,4,1200,No\n",
    "2024-01-23,S001,Laptop,3,3000,No\n",
    "2024-01-23,S001,Phone,8,4000,No\n",
    "2024-01-23,S001,Tablet,8,2400,No\n",
    "2024-01-23,S002,Laptop,6,6000,No\n",
    "2024-01-23,S002,Phone,12,6000,No\n",
    "2024-01-23,S002,Tablet,12,3600,No\n",
    "2024-01-24,S001,Laptop,2,2000,No\n",
    "2024-01-24,S001,Phone,14,7000,No\n",
    "2024-01-24,S001,Tablet,6,1800,No\n",
    "2024-01-24,S002,Laptop,1,1000,No\n",
    "2024-01-24,S002,Phone,5,2500,No\n",
    "2024-01-24,S002,Tablet,7,2100,No\n",
    "2024-01-25,S001,Laptop,7,7000,No\n",
    "2024-01-25,S001,Phone,11,5500,No\n",
    "2024-01-25,S001,Tablet,11,3300,No\n",
    "2024-01-25,S002,Laptop,6,6000,No\n",
    "2024-01-25,S002,Phone,11,5500,No\n",
    "2024-01-25,S002,Tablet,5,1500,No\n",
    "2024-01-26,S001,Laptop,5,5000,Yes\n",
    "2024-01-26,S001,Phone,22,11000,Yes\n",
    "2024-01-26,S001,Tablet,7,2100,Yes\n",
    "2024-01-26,S002,Laptop,6,6000,Yes\n",
    "2024-01-26,S002,Phone,24,12000,Yes\n",
    "2024-01-26,S002,Tablet,3,900,Yes\n",
    "2024-01-27,S001,Laptop,7,7000,Yes\n",
    "2024-01-27,S001,Phone,20,10000,Yes\n",
    "2024-01-27,S001,Tablet,6,1800,Yes\n",
    "2024-01-27,S002,Laptop,4,4000,Yes\n",
    "2024-01-27,S002,Phone,8,4000,Yes\n",
    "2024-01-27,S002,Tablet,6,1800,Yes\n",
    "2024-01-28,S001,Laptop,10,10000,No\n",
    "2024-01-28,S001,Phone,15,7500,No\n",
    "2024-01-28,S001,Tablet,12,3600,No\n",
    "2024-01-28,S002,Laptop,6,6000,No\n",
    "2024-01-28,S002,Phone,11,5500,No\n",
    "2024-01-28,S002,Tablet,10,3000,No\n",
    "2024-01-29,S001,Laptop,3,3000,No\n",
    "2024-01-29,S001,Phone,16,8000,No\n",
    "2024-01-29,S001,Tablet,5,1500,No\n",
    "2024-01-29,S002,Laptop,6,6000,No\n",
    "2024-01-29,S002,Phone,17,8500,No\n",
    "2024-01-29,S002,Tablet,2,600,No\n",
    "2024-01-30,S001,Laptop,3,3000,No\n",
    "2024-01-30,S001,Phone,11,5500,No\n",
    "2024-01-30,S001,Tablet,2,600,No\n",
    "2024-01-30,S002,Laptop,6,6000,No\n",
    "2024-01-30,S002,Phone,16,8000,No\n",
    "2024-01-30,S002,Tablet,8,2400,No\n",
    "2024-01-31,S001,Laptop,5,5000,Yes\n",
    "2024-01-31,S001,Phone,22,11000,Yes\n",
    "2024-01-31,S001,Tablet,9,2700,Yes\n",
    "2024-01-31,S002,Laptop,3,3000,Yes\n",
    "2024-01-31,S002,Phone,14,7000,Yes\n",
    "2024-01-31,S002,Tablet,4,1200,Yes\n",
    "2024-02-01,S001,Laptop,2,2000,No\n",
    "2024-02-01,S001,Phone,7,3500,No\n",
    "2024-02-01,S001,Tablet,11,3300,No\n",
    "2024-02-01,S002,Laptop,6,6000,No\n",
    "2024-02-01,S002,Phone,11,5500,No\n",
    "2024-02-01,S002,Tablet,5,1500,No\n",
    "2024-02-02,S001,Laptop,2,2000,No\n",
    "2024-02-02,S001,Phone,9,4500,No\n",
    "2024-02-02,S001,Tablet,7,2100,No\n",
    "2024-02-02,S002,Laptop,5,5000,No\n",
    "2024-02-02,S002,Phone,9,4500,No\n",
    "2024-02-02,S002,Tablet,12,3600,No\n",
    "2024-02-03,S001,Laptop,9,9000,No\n",
    "2024-02-03,S001,Phone,12,6000,No\n",
    "2024-02-03,S001,Tablet,9,2700,No\n",
    "2024-02-03,S002,Laptop,10,10000,No\n",
    "2024-02-03,S002,Phone,6,3000,No\n",
    "2024-02-03,S002,Tablet,10,3000,No\n",
    "2024-02-04,S001,Laptop,6,6000,No\n",
    "2024-02-04,S001,Phone,5,2500,No\n",
    "2024-02-04,S001,Tablet,8,2400,No\n",
    "2024-02-04,S002,Laptop,6,6000,No\n",
    "2024-02-04,S002,Phone,10,5000,No\n",
    "2024-02-04,S002,Tablet,10,3000,No\n",
    "2024-02-05,S001,Laptop,7,7000,No\n",
    "2024-02-05,S001,Phone,13,6500,No\n",
    "2024-02-05,S001,Tablet,11,3300,No\n",
    "2024-02-05,S002,Laptop,8,8000,No\n",
    "2024-02-05,S002,Phone,11,5500,No\n",
    "2024-02-05,S002,Tablet,8,2400,No\n",
    "2024-02-06,S001,Laptop,5,5000,No\n",
    "2024-02-06,S001,Phone,14,7000,No\n",
    "2024-02-06,S001,Tablet,4,1200,No\n",
    "2024-02-06,S002,Laptop,2,2000,No\n",
    "2024-02-06,S002,Phone,11,5500,No\n",
    "2024-02-06,S002,Tablet,7,2100,No\n",
    "2024-02-07,S001,Laptop,6,6000,No\n",
    "2024-02-07,S001,Phone,7,3500,No\n",
    "2024-02-07,S001,Tablet,9,2700,No\n",
    "2024-02-07,S002,Laptop,2,2000,No\n",
    "2024-02-07,S002,Phone,8,4000,No\n",
    "2024-02-07,S002,Tablet,9,2700,No\n",
    "2024-02-08,S001,Laptop,5,5000,No\n",
    "2024-02-08,S001,Phone,12,6000,No\n",
    "2024-02-08,S001,Tablet,3,900,No\n",
    "2024-02-08,S002,Laptop,8,8000,No\n",
    "2024-02-08,S002,Phone,5,2500,No\n",
    "2024-02-08,S002,Tablet,8,2400,No\n",
    "2024-02-09,S001,Laptop,6,6000,Yes\n",
    "2024-02-09,S001,Phone,18,9000,Yes\n",
    "2024-02-09,S001,Tablet,5,1500,Yes\n",
    "2024-02-09,S002,Laptop,7,7000,Yes\n",
    "2024-02-09,S002,Phone,18,9000,Yes\n",
    "2024-02-09,S002,Tablet,5,1500,Yes\n",
    "2024-02-10,S001,Laptop,9,9000,No\n",
    "2024-02-10,S001,Phone,6,3000,No\n",
    "2024-02-10,S001,Tablet,8,2400,No\n",
    "2024-02-10,S002,Laptop,7,7000,No\n",
    "2024-02-10,S002,Phone,5,2500,No\n",
    "2024-02-10,S002,Tablet,6,1800,No\n",
    "2024-02-11,S001,Laptop,6,6000,No\n",
    "2024-02-11,S001,Phone,11,5500,No\n",
    "2024-02-11,S001,Tablet,2,600,No\n",
    "2024-02-11,S002,Laptop,7,7000,No\n",
    "2024-02-11,S002,Phone,5,2500,No\n",
    "2024-02-11,S002,Tablet,9,2700,No\n",
    "2024-02-12,S001,Laptop,5,5000,No\n",
    "2024-02-12,S001,Phone,5,2500,No\n",
    "2024-02-12,S001,Tablet,4,1200,No\n",
    "2024-02-12,S002,Laptop,1,1000,No\n",
    "2024-02-12,S002,Phone,14,7000,No\n",
    "2024-02-12,S002,Tablet,15,4500,No\n",
    "2024-02-13,S001,Laptop,3,3000,No\n",
    "2024-02-13,S001,Phone,18,9000,No\n",
    "2024-02-13,S001,Tablet,8,2400,No\n",
    "2024-02-13,S002,Laptop,5,5000,No\n",
    "2024-02-13,S002,Phone,8,4000,No\n",
    "2024-02-13,S002,Tablet,6,1800,No\n",
    "2024-02-14,S001,Laptop,4,4000,No\n",
    "2024-02-14,S001,Phone,9,4500,No\n",
    "2024-02-14,S001,Tablet,6,1800,No\n",
    "2024-02-14,S002,Laptop,4,4000,No\n",
    "2024-02-14,S002,Phone,6,3000,No\n",
    "2024-02-14,S002,Tablet,7,2100,No\n",
    "2024-02-15,S001,Laptop,4,4000,Yes\n",
    "2024-02-15,S001,Phone,26,13000,Yes\n",
    "2024-02-15,S001,Tablet,5,1500,Yes\n",
    "2024-02-15,S002,Laptop,2,2000,Yes\n",
    "2024-02-15,S002,Phone,14,7000,Yes\n",
    "2024-02-15,S002,Tablet,6,1800,Yes\n",
    "2024-02-16,S001,Laptop,7,7000,No\n",
    "2024-02-16,S001,Phone,9,4500,No\n",
    "2024-02-16,S001,Tablet,1,300,No\n",
    "2024-02-16,S002,Laptop,6,6000,No\n",
    "2024-02-16,S002,Phone,12,6000,No\n",
    "2024-02-16,S002,Tablet,10,3000,No\n",
    "2024-02-17,S001,Laptop,5,5000,No\n",
    "2024-02-17,S001,Phone,8,4000,No\n",
    "2024-02-17,S001,Tablet,14,4200,No\n",
    "2024-02-17,S002,Laptop,4,4000,No\n",
    "2024-02-17,S002,Phone,13,6500,No\n",
    "2024-02-17,S002,Tablet,7,2100,No\n",
    "2024-02-18,S001,Laptop,6,6000,Yes\n",
    "2024-02-18,S001,Phone,22,11000,Yes\n",
    "2024-02-18,S001,Tablet,9,2700,Yes\n",
    "2024-02-18,S002,Laptop,2,2000,Yes\n",
    "2024-02-18,S002,Phone,10,5000,Yes\n",
    "2024-02-18,S002,Tablet,12,3600,Yes\n",
    "2024-02-19,S001,Laptop,6,6000,No\n",
    "2024-02-19,S001,Phone,12,6000,No\n",
    "2024-02-19,S001,Tablet,3,900,No\n",
    "2024-02-19,S002,Laptop,3,3000,No\n",
    "2024-02-19,S002,Phone,4,2000,No\n",
    "2024-02-19,S002,Tablet,7,2100,No\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell writes the RAG product catalog file, `data/rag/product_catalog.md`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile data/rag/product_catalog.md\n",
    "# Product Catalog: Smartphones, Laptops, and Tablets\n",
    "\n",
    "## Smartphones\n",
    "\n",
    "The Veltrix Solis Z9 is a flagship device in the premium smartphone segment. It builds on a decade of design iterations that prioritize screen-to-body ratio, minimal bezels, and high refresh rate displays. The 6.7-inch AMOLED panel with 120Hz refresh rate delivers immersive visual experiences, whether in gaming, video streaming, or augmented reality applications. The display's GorillaGlass Fusion coating provides scratch resistance and durability, and the thin form factor is engineered using a titanium-aluminum alloy chassis to reduce weight without compromising rigidity.\n",
    "\n",
    "Internally, the Solis Z9 is powered by the OrionEdge V14 chipset, a 4nm process SoC designed for high-efficiency workloads. Its AI accelerator module handles on-device tasks such as voice transcription, camera optimization, and intelligent background app management. The inclusion of 12GB LPDDR5 RAM and a 256GB UFS 3.1 storage system allows for seamless multitasking, instant app launching, and rapid data access. The device supports eSIM and dual physical SIM configurations, catering to global travelers and hybrid network users.\n",
    "\n",
    "Photography and videography are central to the Solis Z9 experience. The triple-camera system incorporates a periscope-style 8MP telephoto lens with 5x optical zoom, a 12MP ultra-wide sensor with macro capabilities, and a 64MP main sensor featuring optical image stabilization (OIS) and phase detection autofocus (PDAF). Night mode and HDRX+ processing enable high-fidelity image capture in challenging lighting conditions.\n",
    "\n",
    "Software-wise, the device ships with LunOS 15, a lightweight Android fork optimized for modular updates and privacy compliance. The system supports secure containers for work profiles and AI-powered notifications that summarize app alerts across channels. Facial unlock is augmented by a 3D IR depth sensor, providing reliable biometric security alongside the ultrasonic in-display fingerprint scanner.\n",
    "\n",
    "The Solis Z9 is a culmination of over a decade of design experimentation in mobile form factors, ranging from curved-edge screens to under-display camera arrays. Its balance of performance, battery efficiency, and user-centric software makes it an ideal daily driver for content creators, mobile gamers, and enterprise users.\n",
    "\n",
    "## Laptops\n",
    "\n",
    "The Cryon Vanta 16X represents the latest evolution of portable computing power tailored for professional-grade workloads.\n",
    "\n",
    "The Vanta 16X features a unibody chassis milled from aircraft-grade aluminum using CNC machining. The thermal design integrates vapor chamber cooling and dual-fan exhaust architecture to support sustained performance under high computational loads. The 16-inch 4K UHD display is color-calibrated at the factory and supports HDR10+, making it suitable for cinematic video editing and high-fidelity CAD modeling.\n",
    "\n",
    "Powering the device is Intel's Core i9-13900H processor, which includes 14 cores with a hybrid architecture combining performance and efficiency cores. This allows the system to dynamically balance power consumption and raw speed based on active workloads. The dedicated Zephira RTX 4700G GPU features 8GB of GDDR6 VRAM and is optimized for CUDA and Tensor Core operations, enabling applications in real-time ray tracing, AI inference, and 3D rendering.\n",
    "\n",
    "The Vanta 16X includes a 2TB PCIe Gen 4 NVMe SSD, delivering sequential read/write speeds above 7GB/s, and 32GB of high-bandwidth DDR5 RAM. The machine supports hardware-accelerated virtualization and dual-booting, and ships with VireoOS Pro pre-installed, with official drivers available for Fedora, Ubuntu LTS, and NebulaOS.\n",
    "\n",
    "Input options are expansive. The keyboard features per-key RGB lighting and programmable macros, while the haptic touchpad supports multi-gesture navigation and palm rejection. Port variety includes dual Thunderbolt 4 ports, a full-size SD Express card reader, HDMI 2.1, 2.5G Ethernet, three USB-A 3.2 ports, and a 3.5mm TRRS audio jack. A fingerprint reader is embedded in the power button and supports biometric logins via Windows Hello.\n",
    "\n",
    "The history of the Cryon laptop line dates back to the early 2010s, when the company launched its first ultrabook aimed at mobile developers. Since then, successive generations have introduced carbon fiber lids, modular SSD bays, and convertible form factors. The Vanta 16X continues this tradition by integrating a customizable BIOS, a modular fan assembly, and a trackpad optimized for creative software like Blender and Adobe Creative Suite.\n",
    "\n",
    "Designed for software engineers, data scientists, film editors, and 3D artists, the Cryon Vanta 16X is a workstation-class laptop in a portable shell.\n",
    "\n",
    "## Tablets\n",
    "\n",
    "The Nebulyn Ark S12 Ultra reflects the current apex of tablet technology, combining high-end hardware with software environments tailored for productivity and creativity.\n",
    "\n",
    "The Ark S12 Ultra is built around a 12.9-inch OLED display that supports 144Hz refresh rate and HDR10+ dynamic range. With a resolution of 2800 x 1752 pixels and a contrast ratio of 1,000,000:1, the screen delivers vibrant color reproduction ideal for design and media consumption. The display supports true tone adaptation and low blue-light filtering for prolonged use.\n",
    "\n",
    "Internally, the tablet uses Qualcomm's Snapdragon 8 Gen 3 SoC, which includes an Adreno 750 GPU and an NPU for on-device AI tasks. The device ships with 16GB LPDDR5X RAM and 512GB of storage with support for NVMe expansion via a proprietary magnetic dock. The 11200mAh battery enables up to 15 hours of typical use and recharges to 80 percent in 45 minutes via 45W USB-C PD.\n",
    "\n",
    "The Ark's history traces back to the original Nebulyn Tab, which launched in 2014 as an e-reader and video streaming device. Since then, the line has evolved through multiple iterations that introduced stylus support, high-refresh screens, and multi-window desktop modes. The current model supports NebulynVerse, a DeX-like environment that allows external display mirroring and full multitasking with overlapping windows and keyboard shortcuts.\n",
    "\n",
    "Input capabilities are central to the Ark S12 Ultra’s appeal. The Pluma Stylus 3 features magnetic charging, 4096 pressure levels, and tilt detection. It integrates haptic feedback to simulate traditional pen strokes and brush textures. The device also supports a SnapCover keyboard that includes a trackpad and programmable shortcut keys. With the stylus and keyboard, users can effectively transform the tablet into a mobile workstation or digital sketchbook.\n",
    "\n",
    "Camera hardware includes a 13MP main sensor and a 12MP ultra-wide front camera with center-stage tracking and biometric unlock. Microphone arrays with beamforming enable studio-quality call audio. Connectivity includes Wi-Fi 7, Bluetooth 5.3, and optional LTE/5G with eSIM.\n",
    "\n",
    "Software support is robust. The device runs NebulynOS 6.0, based on Android 14L, and supports app sandboxing, multi-user profiles, and remote device management. Integration with cloud services, including SketchNimbus and ThoughtSpace, allows for real-time collaboration and syncing of content across devices.\n",
    "\n",
    "This tablet is targeted at professionals who require a balance between media consumption, creativity, and light productivity. Typical users include architects, consultants, university students, and UX designers.\n",
    "\n",
    "## Comparative Summary\n",
    "\n",
    "Each of these devices—the Veltrix Solis Z9, Cryon Vanta 16X, and Nebulyn Ark S12 Ultra—represents a best-in-class interpretation of its category. The Solis Z9 excels in mobile photography and everyday communication. The Vanta 16X is tailored for high-performance applications such as video production and AI prototyping. The Ark S12 Ultra provides a canvas for creativity, note-taking, and hybrid productivity use cases.\n",
    "\n",
    "## Historical Trends and Design Evolution\n",
    "\n",
    "Design across all three categories is converging toward modularity, longevity, and environmental sustainability. Recycled materials, reparability scores, and software longevity are becoming integral to brand reputation and product longevity. Future iterations are expected to feature tighter integration with wearable devices, ambient AI experiences, and cross-device workflows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"installing-nat\"></a>\n",
    "## 0.4) Installing NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The recommended way to install NAT is through `pip` or `uv pip`.\n",
    "\n",
    "First, we will install `uv` which offers parallel downloads and faster dependency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install uv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NeMo Agent toolkit can be installed through the PyPI `nvidia-nat` package.\n",
    "\n",
    "There are several optional subpackages available for NAT. For this example, we will rely on two subpackages:\n",
    "* The `langchain` subpackage contains useful components for integrating and running within [LangChain](https://python.langchain.com/docs/introduction/).\n",
    "* The `llama-index` subpackage contains useful components for integrating and running within [LlamaIndex](https://developers.llamaindex.ai/python/framework/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"nvidia-nat-langchain\"\n",
    "nat_langchain_installed=$?\n",
    "uv pip show -q \"nvidia-nat-llama-index\"\n",
    "nat_llama_index_installed=$?\n",
    "uv pip show -q \"matplotlib\"\n",
    "matplotlib_installed=$?\n",
    "if [[ ${nat_langchain_installed} -ne 0 || ${nat_llama_index_installed} -ne 0 ]]; then\n",
    "    uv pip install \"nvidia-nat[langchain,llama-index]\"\n",
    "else\n",
    "    echo \"nvidia-nat[langchain,llama-index] is already installed\"\n",
    "fi\n",
    "\n",
    "if [[ ${matplotlib_installed} -ne 0 ]]; then\n",
    "    uv pip install \"matplotlib~=3.9\"\n",
    "else\n",
    "    echo \"matplotlib is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"creating-workflow\"></a>\n",
    "# 1.0) Defining Tools for the Multi-Agent Workflow\n",
    "\n",
    "As explained in detail in previous notebooks in this series, we can use the `nat workflow create` sub-command to create the necessary directory structure for a new agent.\n",
    "\n",
    "Within this directory we can define all of the functions that we want to be available to the agent at runtime. In this notebook specifically we are going to extend on the tool calling agent demonstrated in the previous notebook and add a sub-agent to the available tools. We will also make the sub-agent 'discoverable' by the orchestrator by defining it in the workflow's `register.py` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create retail_sales_agent_nb5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next cells we are going to redefine tool calls explained in the previous notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"product-sales-tool\"></a>\n",
    "## 1.1) Total Product Sales Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/total_product_sales_data_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class GetTotalProductSalesDataConfig(FunctionBaseConfig, name=\"get_total_product_sales_data\"):\n",
    "    \"\"\"Get total sales data by product.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=GetTotalProductSalesDataConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def get_total_product_sales_data_function(config: GetTotalProductSalesDataConfig, _builder: Builder):\n",
    "    \"\"\"Get total sales data for a specific product.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _get_total_product_sales_data(product_name: str) -> str:\n",
    "        \"\"\"\n",
    "        Retrieve total sales data for a specific product.\n",
    "\n",
    "        Args:\n",
    "            product_name: Name of the product\n",
    "\n",
    "        Returns:\n",
    "            String message containing total sales data\n",
    "        \"\"\"\n",
    "        df['Product'] = df[\"Product\"].apply(lambda x: x.lower())\n",
    "        revenue = df[df['Product'] == product_name]['Revenue'].sum()\n",
    "        units_sold = df[df['Product'] == product_name]['UnitsSold'].sum()\n",
    "\n",
    "        return f\"Revenue for {product_name} are {revenue} and total units sold are {units_sold}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _get_total_product_sales_data,\n",
    "        description=_get_total_product_sales_data.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"sales-per-day-tool\"></a>\n",
    "## 1.2) Sales Per Day Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/sales_per_day_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class GetSalesPerDayConfig(FunctionBaseConfig, name=\"get_sales_per_day\"):\n",
    "    \"\"\"Get total sales across all products per day.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=GetSalesPerDayConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def sales_per_day_function(config: GetSalesPerDayConfig, builder: Builder):\n",
    "    \"\"\"Get total sales across all products per day.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "    df['Product'] = df[\"Product\"].apply(lambda x: x.lower())\n",
    "\n",
    "    async def _get_sales_per_day(date: str, product: str) -> str:\n",
    "        \"\"\"\n",
    "        Calculate total sales data across all products for a specific date.\n",
    "\n",
    "        Args:\n",
    "            date: Date in YYYY-MM-DD format\n",
    "            product: Product name\n",
    "\n",
    "        Returns:\n",
    "            String message with the total sales for the day\n",
    "        \"\"\"\n",
    "        if date == \"None\":\n",
    "            return \"Please provide a date in YYYY-MM-DD format.\"\n",
    "        total_revenue = df[(df['Date'] == date) & (df['Product'] == product)]['Revenue'].sum()\n",
    "        total_units_sold = df[(df['Date'] == date) & (df['Product'] == product)]['UnitsSold'].sum()\n",
    "\n",
    "        return f\"Total revenue for {date} is {total_revenue} and total units sold is {total_units_sold}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _get_sales_per_day,\n",
    "        description=_get_sales_per_day.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"detect-outliers-tool\"></a>\n",
    "## 1.3) Detect Outliers Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/detect_outliers_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class DetectOutliersIQRConfig(FunctionBaseConfig, name=\"detect_outliers_iqr\"):\n",
    "    \"\"\"Detect outliers in sales data using IQR method.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=DetectOutliersIQRConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def detect_outliers_iqr_function(config: DetectOutliersIQRConfig, _builder: Builder):\n",
    "    \"\"\"Detect outliers in sales data using the Interquartile Range (IQR) method.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _detect_outliers_iqr(metric: str) -> str:\n",
    "        \"\"\"\n",
    "        Detect outliers in retail data using the IQR method.\n",
    "\n",
    "        Args:\n",
    "            metric: Specific metric to check for outliers\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing outlier analysis results\n",
    "        \"\"\"\n",
    "        if metric == \"None\":\n",
    "            column = \"Revenue\"\n",
    "        else:\n",
    "            column = metric\n",
    "\n",
    "        q1 = df[column].quantile(0.25)\n",
    "        q3 = df[column].quantile(0.75)\n",
    "        iqr = q3 - q1\n",
    "        outliers = df[(df[column] < q1 - 1.5 * iqr) | (df[column] > q3 + 1.5 * iqr)]\n",
    "\n",
    "        return f\"Outliers in {column} are {outliers.to_dict('records')}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _detect_outliers_iqr,\n",
    "        description=_detect_outliers_iqr.__doc__)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"technical-specs-tool\"></a>\n",
    "## 1.4) Technical Specs Retrieval Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/retail_sales_rag_tool.py\n",
    "import logging\n",
    "import os\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import EmbedderRef\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class LlamaIndexRAGConfig(FunctionBaseConfig, name=\"retail_sales_rag\"):\n",
    "\n",
    "    llm_name: LLMRef = Field(description=\"The name of the LLM to use for the RAG engine.\")\n",
    "    embedder_name: EmbedderRef = Field(description=\"The name of the embedder to use for the RAG engine.\")\n",
    "    data_dir: str = Field(description=\"The directory containing the data to use for the RAG engine.\")\n",
    "    description: str = Field(description=\"A description of the knowledge included in the RAG system.\")\n",
    "    collection_name: str = Field(default=\"context\", description=\"The name of the collection to use for the RAG engine.\")\n",
    "\n",
    "\n",
    "def _walk_directory(root: str):\n",
    "    for root, dirs, files in os.walk(root):\n",
    "        for file_name in files:\n",
    "            yield os.path.join(root, file_name)\n",
    "\n",
    "\n",
    "@register_function(config_type=LlamaIndexRAGConfig, framework_wrappers=[LLMFrameworkEnum.LLAMA_INDEX])\n",
    "async def retail_sales_rag_tool(config: LlamaIndexRAGConfig, builder: Builder):\n",
    "    from llama_index.core import Settings\n",
    "    from llama_index.core import SimpleDirectoryReader\n",
    "    from llama_index.core import StorageContext\n",
    "    from llama_index.core import VectorStoreIndex\n",
    "    from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)\n",
    "    embedder = await builder.get_embedder(config.embedder_name, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)\n",
    "\n",
    "    Settings.embed_model = embedder\n",
    "    Settings.llm = llm\n",
    "\n",
    "    files = list(_walk_directory(config.data_dir))\n",
    "    docs = SimpleDirectoryReader(input_files=files).load_data()\n",
    "    logger.info(\"Loaded %s documents from %s\", len(docs), config.data_dir)\n",
    "\n",
    "    parser = SentenceSplitter(\n",
    "        chunk_size=400,\n",
    "        chunk_overlap=20,\n",
    "        separator=\" \",\n",
    "    )\n",
    "    nodes = parser.get_nodes_from_documents(docs)\n",
    "\n",
    "    index = VectorStoreIndex(nodes)\n",
    "\n",
    "    query_engine = index.as_query_engine(similarity_top_k=3, )\n",
    "\n",
    "    async def _arun(inputs: str) -> str:\n",
    "        \"\"\"\n",
    "        Search product catalog for information about tablets, laptops, and smartphones\n",
    "        Args:\n",
    "            inputs: user query about product specifications\n",
    "        \"\"\"\n",
    "        try:\n",
    "            response = query_engine.query(inputs)\n",
    "            return str(response.response)\n",
    "\n",
    "        except Exception as e:\n",
    "            logger.error(\"RAG query failed: %s\", e)\n",
    "            return f\"Sorry, I couldn't retrieve information about that product. Error: {str(e)}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(_arun, description=config.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"plotting-tools\"></a>\n",
    "## 1.5) Data Analysis/Plotting Tools\n",
    "\n",
    "This is a new set of tools that will be registered to the data analysis and plotting agent. This set of tools allows the registered agent to plot the results of upstream data analysis tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/data_visualization_tools.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class PlotSalesTrendForStoresConfig(FunctionBaseConfig, name=\"plot_sales_trend_for_stores\"):\n",
    "    \"\"\"Plot sales trend for a specific store.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=PlotSalesTrendForStoresConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def plot_sales_trend_for_stores_function(config: PlotSalesTrendForStoresConfig, _builder: Builder):\n",
    "    \"\"\"Create a visualization of sales trends over time.\"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _plot_sales_trend_for_stores(store_id: str) -> str:\n",
    "        if store_id not in df[\"StoreID\"].unique():\n",
    "            data = df\n",
    "            title = \"Sales Trend for All Stores\"\n",
    "        else:\n",
    "            data = df[df[\"StoreID\"] == store_id]\n",
    "            title = f\"Sales Trend for Store {store_id}\"\n",
    "\n",
    "        plt.figure(figsize=(10, 5))\n",
    "        trend = data.groupby(\"Date\")[\"Revenue\"].sum()\n",
    "        trend.plot(title=title)\n",
    "        plt.xlabel(\"Date\")\n",
    "        plt.ylabel(\"Revenue\")\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(\"sales_trend.png\")\n",
    "\n",
    "        return \"Sales trend plot saved to sales_trend.png\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _plot_sales_trend_for_stores,\n",
    "        description=(\n",
    "            \"This tool can be used to plot the sales trend for a specific store or all stores. \"\n",
    "            \"It takes in a store ID creates and saves an image of a plot of the revenue trend for that store.\"))\n",
    "\n",
    "\n",
    "class PlotAndCompareRevenueAcrossStoresConfig(FunctionBaseConfig, name=\"plot_and_compare_revenue_across_stores\"):\n",
    "    \"\"\"Plot and compare revenue across stores.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=PlotAndCompareRevenueAcrossStoresConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def plot_revenue_across_stores_function(config: PlotAndCompareRevenueAcrossStoresConfig, _builder: Builder):\n",
    "    \"\"\"Create a visualization comparing sales trends between stores.\"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _plot_revenue_across_stores(arg: str) -> str:\n",
    "        pivot = df.pivot_table(index=\"Date\", columns=\"StoreID\", values=\"Revenue\", aggfunc=\"sum\")\n",
    "        pivot.plot(figsize=(12, 6), title=\"Revenue Trends Across Stores\")\n",
    "        plt.xlabel(\"Date\")\n",
    "        plt.ylabel(\"Revenue\")\n",
    "        plt.legend(title=\"StoreID\")\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(\"revenue_across_stores.png\")\n",
    "\n",
    "        return \"Revenue trends across stores plot saved to revenue_across_stores.png\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _plot_revenue_across_stores,\n",
    "        description=(\n",
    "            \"This tool can be used to plot and compare the revenue trends across stores. Use this tool only if the \"\n",
    "            \"user asks for a comparison of revenue trends across stores.\"\n",
    "            \"It takes in a single string as input (which is ignored) and creates and saves an image of a plot of the revenue trends across stores.\"\n",
    "        ))\n",
    "\n",
    "\n",
    "class PlotAverageDailyRevenueConfig(FunctionBaseConfig, name=\"plot_average_daily_revenue\"):\n",
    "    \"\"\"Plot average daily revenue for stores and products.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=PlotAverageDailyRevenueConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def plot_average_daily_revenue_function(config: PlotAverageDailyRevenueConfig, _builder: Builder):\n",
    "    \"\"\"Create a bar chart showing average daily revenue by day of week.\"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _plot_average_daily_revenue(arg: str) -> str:\n",
    "        daily_revenue = df.groupby([\"StoreID\", \"Product\", \"Date\"])[\"Revenue\"].sum().reset_index()\n",
    "\n",
    "        avg_daily_revenue = daily_revenue.groupby([\"StoreID\", \"Product\"])[\"Revenue\"].mean().unstack()\n",
    "\n",
    "        avg_daily_revenue.plot(kind=\"bar\", figsize=(12, 6), title=\"Average Daily Revenue per Store by Product\")\n",
    "        plt.ylabel(\"Average Revenue\")\n",
    "        plt.xlabel(\"Store ID\")\n",
    "        plt.xticks(rotation=0)\n",
    "        plt.legend(title=\"Product\", bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(\"average_daily_revenue.png\")\n",
    "\n",
    "        return \"Average daily revenue plot saved to average_daily_revenue.png\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _plot_average_daily_revenue,\n",
    "        description=(\"This tool can be used to plot the average daily revenue for stores and products \"\n",
    "                     \"It takes in a single string as input and creates and saves an image of a grouped bar chart \"\n",
    "                     \"of the average daily revenue\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"register-tools\"></a>\n",
    "## 1.6) Register The Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb5/src/retail_sales_agent_nb5/register.py\n",
    "\n",
    "from . import sales_per_day_tool\n",
    "from . import detect_outliers_tool\n",
    "from . import total_product_sales_data_tool\n",
    "from . import retail_sales_rag_tool\n",
    "from . import data_visualization_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"adding-orchestrator\"></a>\n",
    "# 2.0) Adding an Agent Orchestrator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"orchestrator-config\"></a>\n",
    "## 2.1) Agent Orchestrator Workflow Configuration File"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we introduce a new workflow configuration file. Upon first glance, this configuration file may seem complex. However, when we classify the available tools and delegate them to an agent, it will begin to make intuitive sense how sub-tasks in a workflow can be divided amongst the best fit agent.\n",
    "\n",
    "In this multi-agent orchestration system, we divide responsibilities among specialized agents, each equipped with tools that match their domain expertise:\n",
    "\n",
    "**1. Data Analysis Agent**\n",
    "- **Tools:** `total_product_sales_data`, `sales_per_day`, `detect_outliers`\n",
    "- **Justification:** This agent handles raw data processing and statistical analysis. These tools extract, aggregate, and analyze sales data, making it the expert for answering questions about sales trends, patterns, and anomalies. By isolating data analysis tasks, we ensure consistent and reliable data interpretation.\n",
    "\n",
    "**2. Visualization Agent**\n",
    "- **Tools:** `plot_total_product_sales`, `plot_sales_per_day`, `plot_average_daily_revenue`\n",
    "- **Justification:** This agent specializes in creating visual representations of data. Visualization requires different expertise than raw data analysis—it involves understanding chart types, formatting, and visual communication. Separating this from data analysis allows the agent to focus on producing clear, effective visualizations without mixing concerns.\n",
    "\n",
    "**3. Knowledge Retrieval Agent (RAG)**\n",
    "- **Tools:** `retail_sales_rag`\n",
    "- **Justification:** This agent accesses external knowledge bases and documentation through retrieval-augmented generation. It handles questions that require contextual information beyond the sales data itself, such as business policies, product information, or historical context. This separation ensures that knowledge retrieval doesn't interfere with computational tasks.\n",
    "\n",
    "**4. Orchestrator Agent (Top-Level)**\n",
    "- **Tools:** None (delegates to sub-agents)\n",
    "- **Justification:** The orchestrator doesn't perform tasks directly but instead routes requests to the appropriate expert agent. This design pattern mirrors real-world organizational structures where a manager delegates to specialists. It enables complex workflows where multiple agents collaborate, each contributing their expertise to solve multi-faceted problems.\n",
    "\n",
    "This architecture provides several benefits:\n",
    "- **Modularity:** Each agent can be updated or replaced independently\n",
    "- **Clarity:** Tool responsibilities are clearly defined and scoped\n",
    "- **Scalability:** New agents and tools can be added without disrupting existing ones\n",
    "- **Efficiency:** Agents only load and reason about tools relevant to their domain\n",
    "\n",
    "> **Note:** _You will notice in the below configuration that no tools are directly called by the workflow-level agent. Instead, it delegates specifically to expert agents based on the request_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/configs/config_multi_agent.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 2048\n",
    "    context_window: 32768\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "embedders:\n",
    "  nim_embedder:\n",
    "    _type: nim\n",
    "    model_name: nvidia/nv-embedqa-e5-v5\n",
    "    truncate: END\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  total_product_sales_data:\n",
    "    _type: get_total_product_sales_data\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  sales_per_day:\n",
    "    _type: get_sales_per_day\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  detect_outliers:\n",
    "    _type: detect_outliers_iqr\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "  data_analysis_agent:\n",
    "    _type: tool_calling_agent\n",
    "    tool_names:\n",
    "      - total_product_sales_data\n",
    "      - sales_per_day\n",
    "      - detect_outliers\n",
    "    llm_name: nim_llm\n",
    "    max_history: 10\n",
    "    max_iterations: 15\n",
    "    description: |\n",
    "      A helpful assistant that can answer questions about the retail sales CSV data.\n",
    "      Use the tools to answer the questions.\n",
    "      Input is a single string.\n",
    "    verbose: false\n",
    "\n",
    "  product_catalog_rag:\n",
    "    _type: retail_sales_rag\n",
    "    llm_name: nim_llm\n",
    "    embedder_name: nim_embedder\n",
    "    collection_name: product_catalog_rag\n",
    "    data_dir: data/rag/\n",
    "    description: \"Search product catalog for TabZen tablet, AeroBook laptop, NovaPhone specifications\"\n",
    "\n",
    "  rag_agent:\n",
    "    _type: react_agent\n",
    "    llm_name: nim_llm\n",
    "    tool_names: [product_catalog_rag]\n",
    "    max_history: 3\n",
    "    max_iterations: 5\n",
    "    max_retries: 2\n",
    "    description: |\n",
    "      An assistant that can only answer questions about products.\n",
    "      Use the product_catalog_rag tool to answer questions about products.\n",
    "      Do not make up any information.\n",
    "    verbose: false\n",
    "\n",
    "  plot_sales_trend_for_stores:\n",
    "    _type: plot_sales_trend_for_stores\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  plot_and_compare_revenue_across_stores:\n",
    "    _type: plot_and_compare_revenue_across_stores\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  plot_average_daily_revenue:\n",
    "    _type: plot_average_daily_revenue\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "  data_visualization_agent:\n",
    "    _type: react_agent\n",
    "    llm_name: nim_llm\n",
    "    tool_names:\n",
    "      - plot_sales_trend_for_stores\n",
    "      - plot_and_compare_revenue_across_stores\n",
    "      - plot_average_daily_revenue\n",
    "    max_history: 10\n",
    "    max_iterations: 15\n",
    "    description: |\n",
    "      You are a data visualization expert.\n",
    "      You can only create plots and visualizations based on user requests.\n",
    "      Only use available tools to generate plots.\n",
    "      You cannot analyze any data.\n",
    "    verbose: false\n",
    "    handle_parsing_errors: true\n",
    "    max_retries: 2\n",
    "    retry_parsing_errors: true\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names: [data_analysis_agent, data_visualization_agent, rag_agent]\n",
    "  llm_name: nim_llm\n",
    "  verbose: true\n",
    "  handle_parsing_errors: true\n",
    "  max_retries: 2\n",
    "  system_prompt: |\n",
    "    Answer the following questions as best you can.\n",
    "    You may communicate and collaborate with various experts to answer the questions.\n",
    "\n",
    "    {tools}\n",
    "\n",
    "    You may respond in one of two formats.\n",
    "    Use the following format exactly to communicate with an expert:\n",
    "\n",
    "    Question: the input question you must answer\n",
    "    Thought: you should always think about what to do\n",
    "    Action: the action to take, should be one of [{tool_names}]\n",
    "    Action Input: the input to the action (if there is no required input, include \"Action Input: None\")\n",
    "    Observation: wait for the expert to respond, do not assume the expert's response\n",
    "\n",
    "    ... (this Thought/Action/Action Input/Observation can repeat N times.)\n",
    "    Use the following format once you have the final answer:\n",
    "\n",
    "    Thought: I now know the final answer\n",
    "    Final Answer: the final answer to the original input question"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"running-orchestrator-workflow\"></a>\n",
    "## 2.2) Running the Workflow\n",
    "\n",
    "Next we can run the workflow:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file retail_sales_agent_nb5/configs/config_multi_agent.yml \\\n",
    "  --input \"What is the Ark S12 Ultra tablet and what are its specifications?\" \\\n",
    "  --input \"How do laptop sales compare to phone sales?\" \\\n",
    "  --input \"Plot average daily revenue\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If images were generated by tool calls you can view them by running the following code cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "from IPython.display import display\n",
    "\n",
    "display(Image(\"./average_daily_revenue.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"adding-custom-agent\"></a>\n",
    "# 3.0) Adding a Custom Agent\n",
    "\n",
    "Besides using inbuilt agents in the workflows, we can also create custom agents using LangGraph or any other framework and bring them into a workflow. We demonstrate this by swapping out the ReAct agent used by the data visualization expert for a custom agent that has human-in-the-loop capability. The agent will ask the user whether they would like a summary of graph content.\n",
    "\n",
    "This exemplifies how complete agent workflows can be wrapped and used as tools by other agents, enabling complex multi-agent orchestration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"hitl-tool\"></a>\n",
    "## 3.1) Human-in-the-Loop (HITL) Approval Tool\n",
    "\n",
    "The following two cells define the approval tool and its registration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/retail_hitl_approval.py\n",
    "import logging\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.context import Context\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "from nat.data_models.interactive import HumanPromptText\n",
    "from nat.data_models.interactive import InteractionResponse\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class HITLApprovalFnConfig(FunctionBaseConfig, name=\"retail_hitl_approval\"):\n",
    "    \"\"\"\n",
    "    This function is used to get the user's response to the prompt.\n",
    "    It will return True if the user responds with 'yes', otherwise False.\n",
    "    \"\"\"\n",
    "\n",
    "    prompt: str = Field(..., description=\"The prompt to use for the HITL function\")\n",
    "\n",
    "\n",
    "@register_function(config_type=HITLApprovalFnConfig)\n",
    "async def hitl_approval_function(config: HITLApprovalFnConfig, builder: Builder):\n",
    "\n",
    "    import re\n",
    "\n",
    "    prompt = f\"{config.prompt} Please confirm if you would like to proceed. Respond with 'yes' or 'no'.\"\n",
    "\n",
    "    async def _arun(unused: str = \"\") -> bool:\n",
    "\n",
    "        nat_context = Context.get()\n",
    "        user_input_manager = nat_context.user_interaction_manager\n",
    "\n",
    "        human_prompt_text = HumanPromptText(text=prompt, required=True, placeholder=\"<your response here>\")\n",
    "        response: InteractionResponse = await user_input_manager.prompt_user_input(human_prompt_text)\n",
    "        response_str = response.content.text.lower()  # type: ignore\n",
    "        selected_option = re.search(r'\\b(yes)\\b', response_str)\n",
    "\n",
    "        if selected_option:\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    yield FunctionInfo.from_fn(_arun,\n",
    "                               description=(\"This function will be used to get the user's response to the prompt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb5/src/retail_sales_agent_nb5/register.py\n",
    "\n",
    "from . import retail_hitl_approval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"graph-summarizer-tool\"></a>\n",
    "## 3.2) Graph Summarizer Tool\n",
    "\n",
    "The following two cells define the graph summarizer tool and its registration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/graph_summarizer_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class GraphSummarizerConfig(FunctionBaseConfig, name=\"graph_summarizer\"):\n",
    "    \"\"\"Analyze and summarize chart data.\"\"\"\n",
    "    llm_name: LLMRef = Field(description=\"The name of the LLM to use for the graph summarizer.\")\n",
    "\n",
    "\n",
    "@register_function(config_type=GraphSummarizerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def graph_summarizer_function(config: GraphSummarizerConfig, builder: Builder):\n",
    "    \"\"\"Analyze chart data and provide natural language summaries.\"\"\"\n",
    "    import base64\n",
    "\n",
    "    from langchain_core.messages import HumanMessage\n",
    "\n",
    "    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "\n",
    "    async def _graph_summarizer(image_path: str) -> str:\n",
    "        \"\"\"\n",
    "        Analyze chart data and provide insights and summaries.\n",
    "\n",
    "        Args:\n",
    "            image_path: The path to the image to analyze\n",
    "\n",
    "        Returns:\n",
    "            String containing analysis and insights\n",
    "        \"\"\"\n",
    "\n",
    "        def encode_image(image_path: str) -> str:\n",
    "            with open(image_path, \"rb\") as image_file:\n",
    "                return base64.b64encode(image_file.read()).decode('utf-8')\n",
    "\n",
    "        base64_image = encode_image(image_path)\n",
    "\n",
    "        # Create a multimodal message with text and image\n",
    "        message = HumanMessage(\n",
    "            content=[\n",
    "                {\"type\": \"text\", \"text\": \"Please summarize the key insights from this graph in natural language.\"},\n",
    "                {\"type\": \"image_url\", \"image_url\": {\"url\": f\"data:image/png;base64,{base64_image}\"}}\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        # Invoke the LLM with the multimodal message\n",
    "        response = await llm.ainvoke([message])\n",
    "        return response.content\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _graph_summarizer,\n",
    "        description=(\"This tool can be used to summarize the key insights from a graph in natural language. \"\n",
    "                     \"It takes in the path to an image and returns a summary of the key insights from the graph.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb5/src/retail_sales_agent_nb5/register.py\n",
    "\n",
    "from . import graph_summarizer_tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"custom-viz-agent\"></a>\n",
    "## 3.3) Custom Data Visualization Agent With HITL Approval\n",
    "\n",
    "The following two cells define the custom agent and its registration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/src/retail_sales_agent_nb5/data_visualization_agent.py\n",
    "import logging\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function import Function\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import FunctionRef\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class DataVisualizationAgentConfig(FunctionBaseConfig, name=\"data_visualization_agent\"):\n",
    "    \"\"\"\n",
    "    NeMo Agent toolkit function config for data visualization.\n",
    "    \"\"\"\n",
    "    llm_name: LLMRef = Field(description=\"The name of the LLM to use\")\n",
    "    tool_names: list[FunctionRef] = Field(description=\"The names of the tools to use\")\n",
    "    description: str = Field(description=\"The description of the agent.\")\n",
    "    prompt: str = Field(description=\"The prompt to use for the agent.\")\n",
    "    graph_summarizer_fn: FunctionRef = Field(description=\"The function to use for the graph summarizer.\")\n",
    "    hitl_approval_fn: FunctionRef = Field(description=\"The function to use for the hitl approval.\")\n",
    "    max_retries: int = Field(default=3, description=\"The maximum number of retries for the agent.\")\n",
    "\n",
    "\n",
    "@register_function(config_type=DataVisualizationAgentConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def data_visualization_agent_function(config: DataVisualizationAgentConfig, builder: Builder):\n",
    "    from langchain_core.messages import AIMessage\n",
    "    from langchain_core.messages import BaseMessage\n",
    "    from langchain_core.messages import HumanMessage\n",
    "    from langchain_core.messages import SystemMessage\n",
    "    from langchain_core.messages import ToolMessage\n",
    "    from langgraph.graph import StateGraph\n",
    "    from langgraph.prebuilt import ToolNode\n",
    "    from pydantic import BaseModel\n",
    "\n",
    "    class AgentState(BaseModel):\n",
    "        retry_count: int = 0\n",
    "        messages: list[BaseMessage]\n",
    "        approved: bool = True\n",
    "\n",
    "    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "    llm = await builder.get_llm(llm_name=config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)\n",
    "    llm_n_tools = llm.bind_tools(tools)\n",
    "\n",
    "    hitl_approval_fn: Function = await builder.get_function(config.hitl_approval_fn)\n",
    "    graph_summarizer_fn: Function = await builder.get_function(config.graph_summarizer_fn)\n",
    "\n",
    "    async def conditional_edge(state: AgentState):\n",
    "        try:\n",
    "            logger.debug(\"Starting the Tool Calling Conditional Edge\")\n",
    "            messages = state.messages\n",
    "            last_message = messages[-1]\n",
    "            logger.info(\"Last message type: %s\", type(last_message))\n",
    "            logger.info(\"Has tool_calls: %s\", hasattr(last_message, 'tool_calls'))\n",
    "            if hasattr(last_message, 'tool_calls'):\n",
    "                logger.info(\"Tool calls: %s\", last_message.tool_calls)\n",
    "\n",
    "            if (hasattr(last_message, 'tool_calls') and last_message.tool_calls and len(last_message.tool_calls) > 0):\n",
    "                logger.info(\"Routing to tools - found non-empty tool calls\")\n",
    "                return \"tools\"\n",
    "            logger.info(\"Routing to check_hitl_approval - no tool calls to execute\")\n",
    "            return \"check_hitl_approval\"\n",
    "        except Exception as ex:\n",
    "            logger.error(\"Error in conditional_edge: %s\", ex)\n",
    "            if hasattr(state, 'retry_count') and state.retry_count >= config.max_retries:\n",
    "                logger.warning(\"Max retries reached, returning without meaningful output\")\n",
    "                return \"__end__\"\n",
    "            state.retry_count = getattr(state, 'retry_count', 0) + 1\n",
    "            logger.warning(\n",
    "                \"Error in the conditional edge: %s, retrying %d times out of %d\",\n",
    "                ex,\n",
    "                state.retry_count,\n",
    "                config.max_retries,\n",
    "            )\n",
    "            return \"data_visualization_agent\"\n",
    "\n",
    "    def approval_conditional_edge(state: AgentState):\n",
    "        \"\"\"Route to summarizer if user approved, otherwise end\"\"\"\n",
    "        logger.info(\"Approval conditional edge: %s\", state.approved)\n",
    "        if hasattr(state, 'approved') and not state.approved:\n",
    "            return \"__end__\"\n",
    "        return \"summarize\"\n",
    "\n",
    "    def data_visualization_agent(state: AgentState):\n",
    "        sys_msg = SystemMessage(content=config.prompt)\n",
    "        messages = state.messages\n",
    "\n",
    "        if messages and isinstance(messages[-1], ToolMessage):\n",
    "            last_tool_msg = messages[-1]\n",
    "            logger.info(\"Processing tool result: %s\", last_tool_msg.content)\n",
    "            summary_content = f\"I've successfully created the visualization. {last_tool_msg.content}\"\n",
    "            return {\"messages\": [AIMessage(content=summary_content)]}\n",
    "        logger.info(\"Normal agent operation - generating response for: %s\", messages[-1] if messages else 'no messages')\n",
    "        return {\"messages\": [llm_n_tools.invoke([sys_msg] + state.messages)]}\n",
    "\n",
    "    async def check_hitl_approval(state: AgentState):\n",
    "        messages = state.messages\n",
    "        last_message = messages[-1]\n",
    "        logger.info(\"Checking hitl approval: %s\", state.approved)\n",
    "        logger.info(\"Last message type: %s\", type(last_message))\n",
    "        selected_option = await hitl_approval_fn.acall_invoke()\n",
    "        if selected_option:\n",
    "            return {\"approved\": True}\n",
    "        return {\"approved\": False}\n",
    "\n",
    "    async def summarize_graph(state: AgentState):\n",
    "        \"\"\"Summarize the graph using the graph summarizer function\"\"\"\n",
    "        image_path = None\n",
    "        for msg in state.messages:\n",
    "            if hasattr(msg, 'content') and msg.content:\n",
    "                content = str(msg.content)\n",
    "                import re\n",
    "                img_ext = r'[a-zA-Z0-9_.-]+\\.(?:png|jpg|jpeg|gif|svg)'\n",
    "                pattern = rf'saved to ({img_ext})|({img_ext})'\n",
    "                match = re.search(pattern, content)\n",
    "                if match:\n",
    "                    image_path = match.group(1) or match.group(2)\n",
    "                    break\n",
    "\n",
    "        if not image_path:\n",
    "            image_path = \"sales_trend.png\"\n",
    "\n",
    "        logger.info(\"Extracted image path for summarization: %s\", image_path)\n",
    "        response = await graph_summarizer_fn.ainvoke(image_path)\n",
    "        return {\"messages\": [response]}\n",
    "\n",
    "    try:\n",
    "        logger.debug(\"Building and compiling the Agent Graph\")\n",
    "        builder_graph = StateGraph(AgentState)\n",
    "\n",
    "        builder_graph.add_node(\"data_visualization_agent\", data_visualization_agent)\n",
    "        builder_graph.add_node(\"tools\", ToolNode(tools))\n",
    "        builder_graph.add_node(\"check_hitl_approval\", check_hitl_approval)\n",
    "        builder_graph.add_node(\"summarize\", summarize_graph)\n",
    "\n",
    "        builder_graph.add_conditional_edges(\"data_visualization_agent\", conditional_edge)\n",
    "\n",
    "        builder_graph.set_entry_point(\"data_visualization_agent\")\n",
    "        builder_graph.add_edge(\"tools\", \"data_visualization_agent\")\n",
    "\n",
    "        builder_graph.add_conditional_edges(\"check_hitl_approval\", approval_conditional_edge)\n",
    "\n",
    "        builder_graph.add_edge(\"summarize\", \"__end__\")\n",
    "\n",
    "        agent_executor = builder_graph.compile()\n",
    "\n",
    "        logger.info(\"Data Visualization Agent Graph built and compiled successfully\")\n",
    "\n",
    "    except Exception as ex:\n",
    "        logger.error(\"Failed to build Data Visualization Agent Graph: %s\", ex)\n",
    "        raise\n",
    "\n",
    "    async def _arun(user_query: str) -> str:\n",
    "        \"\"\"\n",
    "        Visualize data based on user query.\n",
    "\n",
    "        Args:\n",
    "            user_query (str): User query to visualize data\n",
    "\n",
    "        Returns:\n",
    "            str: Visualization conclusion from the LLM agent\n",
    "        \"\"\"\n",
    "        input_message = f\"User query: {user_query}.\"\n",
    "        response = await agent_executor.ainvoke({\"messages\": [HumanMessage(content=input_message)]})\n",
    "\n",
    "        return response\n",
    "\n",
    "    try:\n",
    "        yield FunctionInfo.from_fn(_arun, description=config.description)\n",
    "\n",
    "    except GeneratorExit:\n",
    "        print(\"Function exited early!\")\n",
    "    finally:\n",
    "        print(\"Cleaning up retail_sales_agent workflow.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb5/src/retail_sales_agent_nb5/register.py\n",
    "\n",
    "from . import data_visualization_agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"custom-agent-config\"></a>\n",
    "## 3.4) Custom Agent Workflow Configuration File\n",
    "\n",
    "Next, we define the workflow configuration file for this custom agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The high-level changes include:\n",
    "- switching from a ReAct agent to the custom agent with HITL\n",
    "- adding additional tools (HITL, graph summarization)\n",
    "- adding an OpenAI LLM for image summarization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb5/configs/config_multi_agent_hitl.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 2048\n",
    "    context_window: 32768\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "  summarizer_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.2-90b-vision-instruct\n",
    "    temperature: 0.0\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "embedders:\n",
    "  nim_embedder:\n",
    "    _type: nim\n",
    "    model_name: nvidia/nv-embedqa-e5-v5\n",
    "    truncate: END\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  total_product_sales_data:\n",
    "    _type: get_total_product_sales_data\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  sales_per_day:\n",
    "    _type: get_sales_per_day\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  detect_outliers:\n",
    "    _type: detect_outliers_iqr\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "  data_analysis_agent:\n",
    "    _type: tool_calling_agent\n",
    "    tool_names:\n",
    "      - total_product_sales_data\n",
    "      - sales_per_day\n",
    "      - detect_outliers\n",
    "    llm_name: nim_llm\n",
    "    max_history: 10\n",
    "    max_iterations: 15\n",
    "    description: |\n",
    "      A helpful assistant that can answer questions about the retail sales CSV data.\n",
    "      Use the tools to answer the questions.\n",
    "      Input is a single string.\n",
    "    verbose: false\n",
    "\n",
    "  plot_sales_trend_for_stores:\n",
    "    _type: plot_sales_trend_for_stores\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  plot_and_compare_revenue_across_stores:\n",
    "    _type: plot_and_compare_revenue_across_stores\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  plot_average_daily_revenue:\n",
    "    _type: plot_average_daily_revenue\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "  retail_hitl_approval:\n",
    "    _type: retail_hitl_approval\n",
    "    prompt: |\n",
    "      Do you want to summarize the created graph content?\n",
    "  graph_summarizer:\n",
    "    _type: graph_summarizer\n",
    "    llm_name: summarizer_llm\n",
    "\n",
    "  data_visualization_agent:\n",
    "    _type: data_visualization_agent\n",
    "    llm_name: nim_llm\n",
    "    tool_names:\n",
    "      - plot_sales_trend_for_stores\n",
    "      - plot_and_compare_revenue_across_stores\n",
    "      - plot_average_daily_revenue\n",
    "    graph_summarizer_fn: graph_summarizer\n",
    "    hitl_approval_fn: retail_hitl_approval\n",
    "    prompt: |\n",
    "      You are a data visualization expert.\n",
    "      Your task is to create plots and visualizations based on user requests.\n",
    "      Use available tools to analyze data and generate plots.\n",
    "    description: |\n",
    "      This is a data visualization agent that should be called if the user asks for a visualization or plot of the data.\n",
    "      It has access to the following tools:\n",
    "      - plot_sales_trend_for_stores: This tool can be used to plot the sales trend for a specific store or all stores.\n",
    "      - plot_and_compare_revenue_across_stores: This tool can be used to plot and compare the revenue trends across stores. Use this tool only if the user asks for a comparison of revenue trends across stores.\n",
    "      - plot_average_daily_revenue: This tool can be used to plot the average daily revenue for stores and products.\n",
    "      The agent will use the available tools to analyze data and generate plots.\n",
    "      The agent will also use the graph_summarizer tool to summarize the graph data.\n",
    "      The agent will also use the retail_hitl_approval to ask the user whether they would like a summary of the graph data.\n",
    "\n",
    "  product_catalog_rag:\n",
    "    _type: retail_sales_rag\n",
    "    llm_name: nim_llm\n",
    "    embedder_name: nim_embedder\n",
    "    collection_name: product_catalog_rag\n",
    "    data_dir: data/rag/\n",
    "    description: \"Search product catalog for TabZen tablet, AeroBook laptop, NovaPhone specifications\"\n",
    "\n",
    "  rag_agent:\n",
    "    _type: react_agent\n",
    "    llm_name: nim_llm\n",
    "    tool_names:\n",
    "      - product_catalog_rag\n",
    "    max_history: 3\n",
    "    max_iterations: 5\n",
    "    max_retries: 2\n",
    "    retry_parsing_errors: true\n",
    "    description: |\n",
    "      An assistant that can answer questions about products.\n",
    "      Use product_catalog_rag to answer questions about products.\n",
    "      Do not make up information.\n",
    "    verbose: true\n",
    "\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names:\n",
    "    - data_analysis_agent\n",
    "    - data_visualization_agent\n",
    "    - rag_agent\n",
    "  llm_name: summarizer_llm\n",
    "  verbose: true\n",
    "  handle_parsing_errors: true\n",
    "  max_retries: 2\n",
    "  system_prompt: |\n",
    "    Answer the following questions as best you can. You may communicate and collaborate with various experts to answer the questions:\n",
    "\n",
    "    {tools}\n",
    "\n",
    "    If the user responds \"no\" to a request to continue, you should end the conversation.\n",
    "\n",
    "    You may respond in one of two formats.\n",
    "    Use the following format exactly to communicate with an expert:\n",
    "\n",
    "    Question: the input question you must answer\n",
    "    Thought: you should always think about what to do\n",
    "    Action: the action to take, should be one of [{tool_names}]\n",
    "    Action Input: the input to the action (if there is no required input, include \"Action Input: None\")\n",
    "    Observation: wait for the expert to respond, do not assume the expert's response\n",
    "\n",
    "    ... (this Thought/Action/Action Input/Observation can repeat N times.)\n",
    "    Use the following format once you have the final answer:\n",
    "\n",
    "    Thought: I now know the final answer\n",
    "    Final Answer: the final answer to the original input question"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"running-custom-workflow\"></a>\n",
    "## 3.5) Running the Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The HITL tool will prompt the user for input when it is called. This works when running in an interactive terminal, but in a notebook environment like this one, the prompt may not function as expected, so we will simulate user input by echoing \"yes\" into the command. This can be replaced with \"no\" to simulate a negative response, or run directly in a terminal for interactive input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! echo \"yes\" | nat run --config_file retail_sales_agent_nb5/configs/config_multi_agent_hitl.yml \\\n",
    "    --input \"Plot average daily revenue\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This concludes this example. We've gone through several examples of integrating tools and custom agents in NeMo Agent toolkit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"next-steps\"></a>\n",
    "# 4.0) Next Steps\n",
    "\n",
    "The next notebook in this series is `observability_evaluation_and_profiling.ipynb` where we will demonstrate how to equip your agent with profiling instrumentation for improved observability over your platform."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
